#include "../DWMRICore/DWI.h"
#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/TwoTensor.h"

#include "TwoTensorFitting.h"

#include <levmar/levmar.h>

#include <teem/ell.h>
#include <teem/ten.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <omp.h>

struct TwoTensor {
	float phiA;
	float phiB;
	CVector3F dirA;
	CVector3F dirB;
	float residual;
};


char *pDWIPathname = NULL;
char *pDTIPathname = NULL;
char *pModelSelectionPathname = NULL;
char *twoTensorPathname = NULL;

/* MLS fiber tracking parameters */
float MLSSize = 4.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;

int w = 0;
int h = 0;
int d = 0;
int	n = 0;
int	m = 0;
float bValue = 0.0f;

int *pIndex = NULL;
CPoint3I *pTwoTensorPos = NULL;
CDWI *pDWI = NULL;
CDTI *pDTI = NULL;
CScalar *pModelSelection = NULL;
CVector3F *pEigenVector = NULL;

int L = 0;
float *phi = NULL;
TwoTensor **pTwoTensorBuffer = NULL;



void PrintUsage()
{
	printf("Usage: TwoTensorFitting -dwi dwi_file -dti dti_file -modelselection model_selection_file -dst dst_file [-thresh threshold]\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dwi") == 0) {
			pDWIPathname = argv[++i];
		} else if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-modelselection") == 0) {
			pModelSelectionPathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			twoTensorPathname = argv[++i];
		} else {
			return -1;
		}
	}

	if (pDWIPathname == NULL || pDTIPathname == NULL || 
		pModelSelectionPathname == NULL || twoTensorPathname == NULL)
		return -2;

	return 0;
}

void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

void DTIMLSFilter(const CPoint3F pos, float *lastDTI, float *filteredDTI)
{
	float mat[9];
	TEN_T2M(mat, lastDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * mat[0] + pSamplePoints[i].m_y * mat[1] + pSamplePoints[i].m_z * mat[2];
		samples[i].m_y = pSamplePoints[i].m_x * mat[3] + pSamplePoints[i].m_y * mat[4] + pSamplePoints[i].m_z * mat[5];
		samples[i].m_z = pSamplePoints[i].m_x * mat[6] + pSamplePoints[i].m_y * mat[7] + pSamplePoints[i].m_z * mat[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	memset(filteredDTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pModelSelection->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			filteredDTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		filteredDTI[j] = filteredDTI[j] / weight;
	}
}

void DWIMLSFilter(const CPoint3F pos, float *filteredDTI, float *filteredDWI, float *filteredS0)
{
	float mat[9];
	TEN_T2M(mat, filteredDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * mat[0] + pSamplePoints[i].m_y * mat[1] + pSamplePoints[i].m_z * mat[2];
		samples[i].m_y = pSamplePoints[i].m_x * mat[3] + pSamplePoints[i].m_y * mat[4] + pSamplePoints[i].m_z * mat[5];
		samples[i].m_z = pSamplePoints[i].m_x * mat[6] + pSamplePoints[i].m_y * mat[7] + pSamplePoints[i].m_z * mat[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	uint n = pDWI->GetN() - pDWI->GetM();
	memset(filteredDWI, 0, sizeof(float)*n);
	*filteredS0 = 0.0f;
	float *temp_DWI = new float[n];
	float temp_S0 = 0.0f;
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pModelSelection->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDWI->GetNS0(samples[i], temp_DWI);
		temp_S0 = pDWI->GetS0(samples[i]);

		for (uint j = 0; j < n; ++j) {
			filteredDWI[j] += pSamplePointsWeight[i] * temp_DWI[j];
		}
		*filteredS0 += pSamplePointsWeight[i] * temp_S0;
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < n; ++j) {
		filteredDWI[j] = filteredDWI[j] / weight;
	}
	*filteredS0 = *filteredS0 / weight;

	SafeDeleteArray(temp_DWI);
}


int CountTwoTensorVoxels()
{
	pIndex = new int[w*h*d];
	memset(pIndex, 0, sizeof(int)*w*h*d);

	int count = 0;
	for (int i = 0; i < w*h*d; ++i) {
		if (pModelSelection->GetScalarFast(i) < 1.5f) {
			pIndex[i] = -1;
		} else {
			pIndex[i] = count;
			count++;
		}
	}

	pTwoTensorPos = new CPoint3I[count];
	for (int z = 0; z < d; ++z) {
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				int i = (z * h + y) * w + x;
				if (pIndex[i] == -1)
					continue;

				int id = pIndex[i];
				pTwoTensorPos[id] = CPoint3I(x, y, z);
			}
		}
	}

	return count;
}

void ComputeEigenVector()
{
	pEigenVector = new CVector3F[w*h*d];
	memset(pEigenVector, 0, sizeof(CVector3F)*w*h*d);

	for (int z = 0; z < d; ++z) {
//	omp_set_num_threads(16);
//#pragma omp parallel for
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				CPoint3F pos = CPoint3F(x, y, z);
				if (pModelSelection->GetScalarFast(x, y, z) < 0.5f) {
					pEigenVector[(z*h+y)*w+x] = CVector3F(0.0f, 0.0f, 0.0f);
				} else {
					float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
					float filteredDTI[7];
					DTIMLSFilter(pos, initDTI, filteredDTI);
					CVector3F dir;
					pDTI->DTI2Direction(filteredDTI, &dir);
					pEigenVector[(z*h+y)*w+x] = dir;
				}
			}
		}
	}
}

void ComputeResidual(const int count)
{
	pTwoTensorBuffer = new TwoTensor*[count];
	for (int i = 0; i < count; ++i) {
		pTwoTensorBuffer[i] = new TwoTensor[L];
		for (int l = 0; l < L; ++l) {
			pTwoTensorBuffer[i][l].phiA = phi[l*2+0];
			pTwoTensorBuffer[i][l].phiB = phi[l*2+1];
		}
	}

	CVector3F *diffusionGradients = pDWI->GetDiffusionGradients();
		
	omp_set_num_threads(16);
#pragma omp parallel for
	for (int i = 0; i < count; ++i) {
		CPoint3F pos;
		pos.m_x = pTwoTensorPos[i].x;
		pos.m_y = pTwoTensorPos[i].y;
		pos.m_z = pTwoTensorPos[i].z;

		float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
		float filteredDTI[7];
		DTIMLSFilter(pos, initDTI, filteredDTI);

		float *filteredDWI = new float[n-m];
		float filteredS0 = 0.0f;
		DWIMLSFilter(pos, filteredDTI, filteredDWI, &filteredS0);

		float eval[3], evec[9];
		pDTI->DTI2Eigen(filteredDTI, eval, evec);

		// rotate the graidents
		CVector3F *gradients = new CVector3F[n-m];
		for (int k = m; k < n; ++k) {
			float v[3], vv[3];
			v[0] = diffusionGradients[k].m_x;
			v[1] = diffusionGradients[k].m_y;
			v[2] = diffusionGradients[k].m_z;
			ELL_3MV_MUL(vv, evec, v);
			gradients[k-m].m_x = vv[0];
			gradients[k-m].m_y = vv[1];
			gradients[k-m].m_z = vv[2];

			filteredDWI[k-m] = filteredDWI[k-m] / filteredS0;
		}

		for (int l = 0; l < L; ++l) {
			float lamda = eval[0];
			//pTwoTensorBuffer[i][l].residual = ConmputeTwoTensorFittingResidual(lamda, filteredDWI,
			//	gradients, n-m, bValue, eval, evec, pTwoTensorBuffer[i][l].phiA, pTwoTensorBuffer[i][l].phiB);
			pTwoTensorBuffer[i][l].dirA = ComputeDirection(eval, evec, lamda, pTwoTensorBuffer[i][l].phiA);
			pTwoTensorBuffer[i][l].dirB = ComputeDirection(eval, evec, lamda, pTwoTensorBuffer[i][l].phiB);
		}

		SafeDeleteArray(filteredDWI);
		SafeDeleteArray(gradients);
	}
}

void ComputePhi()
{
	int N = 18;
	L = N * (N - 1) / 2;

	float stepSize = (float)M_PI / (float)N;

	phi = new float[L*2];
	int id = 0;
	for (int i = 0; i < N; ++i) {
		for (int j = i+1; j < N; ++j) {
			phi[id*2+0] = i * stepSize;
			phi[id*2+1] = j * stepSize;
			id++;
		}
	}
}

float ComputeDP(const int i, const int l, float *deltaPhi, CVector3F *v)
{
	CPoint3I pos = pTwoTensorPos[i];

	int min_x = pos.x - 1;
	int max_x = pos.x + 1;
	int min_y = pos.y - 1;
	int max_y = pos.y + 1;
	int min_z = pos.z - 1;
	int max_z = pos.z + 1;

	float dp = 0.0f;
	for (int z = min_z; z <= max_z; ++z) {
		for (int y = min_y; y <= max_y; ++y) {
			for (int x = min_x; x <= max_x; ++x) {
				if (x == pos.x && y == pos.y && z == pos.z)
					continue;

				float dx = (float)x - pos.x;
				float dy = (float)y - pos.y;
				float dz = (float)z - pos.z;
				float dij = sqrtf(dx * dx + dy * dy + dz * dz);
				float Iij = 0.0f;
				if (pModelSelection->GetScalarFast(x, y, z) > 1.5f ) {
					// planar neighbour
					int j = pIndex[(z*h+y)*w+x];
					float delta_phi_l = pTwoTensorBuffer[i][l].phiA - pTwoTensorBuffer[i][l].phiB;
					float delta_phi_j = deltaPhi[j];
					CVector3F vil = pTwoTensorBuffer[i][l].dirA + pTwoTensorBuffer[i][l].dirB;
					vil.Normalize();
					CVector3F vj = v[j];
					float p1 = (delta_phi_l - delta_phi_j)  / 0.349066f;
					float p2 = acosf(InnerProduct(vil, vj)) / 0.174533f;

					Iij = exp(-p1*p1-p2*p2);
				} else {
					// non-planar neighbour
					CVector3F ej = pEigenVector[(z*h+y)*w+x];
					float p1 = acos(InnerProduct(ej, pTwoTensorBuffer[i][l].dirA)) / 0.174533f;
					float p2 = acos(InnerProduct(ej, pTwoTensorBuffer[i][l].dirB)) / 0.174533f;
					float p = p1 < p2 ? p1 : p2;
					
					Iij = exp(-p*p);
				}

				dp += Iij / dij;
			}
		}
	}
	return dp;
}

void ComputeDeltaPhiAndV(float *deltaPhi, CVector3F *v, float *p, const int count)
{
	for (int i = 0; i < count; ++i) {
		float maxP = 0.0f;
		for (int l = 0; l < L; ++l) {
			if (p[i*L+l] > maxP) {
				maxP = p[i*L+l];
				deltaPhi[i] = pTwoTensorBuffer[i][l].phiA - pTwoTensorBuffer[i][l].phiB;
				v[i] = pTwoTensorBuffer[i][l].dirA + pTwoTensorBuffer[i][l].dirB;
			}
		}
		v[i].Normalize();
	}
}

void SaveResults(float *p)
{
	CVector3F *results = new CVector3F[w*h*d*2];
	for (int i = 0; i < w*h*d; ++i) {
		if (pModelSelection->GetScalarFast(i) > 1.5f) {
			int id = pIndex[i];
			float max_p = 0.0f;
			int max_l = 0;
			for (int l = 0; l < L; ++l) {
				if (p[id*L+l] > max_p) {
					max_p = p[id*L+l];
					max_l = l;
				}
			}
			results[i*2+0] = pTwoTensorBuffer[id][max_l].dirA;
			results[i*2+1] = pTwoTensorBuffer[id][max_l].dirB;
		} else {
			results[i*2+0] = pEigenVector[i];
			results[i*2+1] = CVector3F(0.0f, 0.0f, 0.0f);
		}
	}
	FILE *fp = fopen(twoTensorPathname, "wb");
	fwrite(results, sizeof(CVector3F), w*h*d*2, fp);
	fclose(fp);
	SafeDeleteArray(results);
}

void RegularizedTwoTensorFitting()
{
	w = pDWI->GetWidth();
	h = pDWI->GetHeight();
	d = pDWI->GetDepth();

	n = pDWI->GetN();
	m = pDWI->GetM();
	bValue = pDWI->GetBValue();

	ComputeEigenVector();

	ComputePhi();

	int count = CountTwoTensorVoxels();

	ComputeResidual(count);

	float *residual = new float[count*L];
	FILE *fp = fopen("C:\\Ziang\\Dataset\\Alex\\2X\\ForFiberTracking\\residual.raw", "rb");
	fread(residual, sizeof(float), count*L, fp);
	fclose(fp);
	for (int i = 0; i < count; ++i) {
		for (int l = 0; l < L; ++l) {
			pTwoTensorBuffer[i][l].residual = residual[i*L+l]; 
		}
	}
	SafeDeleteArray(residual);
	
	//exit(1);

	// compute p
	float *p = new float[count*L];
	for (int i = 0; i < count; ++i) {
		float T = 0.0f;
		for (int l = 0; l < L; ++l) {
			T += 1.0f / pTwoTensorBuffer[i][l].residual;
		}
		for (int l = 0; l < L; ++l) {
			p[i*L+l] = 1.0f / (T * pTwoTensorBuffer[i][l].residual);
		}
	}

	// relaxation labeling
	while (true) {
		float *deltaPhi = new float[count];
		CVector3F *v = new CVector3F[count];
		ComputeDeltaPhiAndV(deltaPhi, v, p, count);
		float *nextP = new float[count*L];
//	omp_set_num_threads(16);
//#pragma omp parallel for
		for (int i = 0; i < count; ++i) {
			float *dp = new float[L];
			for (int l = 0; l < L; ++l) {
				dp[l] = ComputeDP(i, l, deltaPhi, v);
			}
			float temp = 0.0f;
			for (int l = 0; l < L; ++l) {
				temp += p[i*L+l] * (1.0f + dp[l]);
			}
			for (int l = 0; l < L; ++l) {
				nextP[i*L+l] = p[i*L+l] * (1.0f + dp[l]) / temp;
			}
			SafeDeleteArray(dp);
		}

		float average_delta = 0.0f;
		for (int i = 0; i < count; ++i) {
			for (int l = 0; l < L; ++l) {
				average_delta += fabs(nextP[i*L+l] - p[i*L+l]);
			}
		}
		average_delta = average_delta / (float)(count*L);

		memcpy(p, nextP, sizeof(float)*count*L);

		SafeDeleteArray(deltaPhi);
		SafeDeleteArray(v);
		SafeDeleteArray(nextP);

		printf("average delta = %f\n", average_delta);

		if (average_delta < 0.0001f)
			break;
	}

	// get results
	SaveResults(p);

	// free all memory
	SafeDeleteArray(pEigenVector);
	SafeDeleteArray(pIndex);
	SafeDeleteArray(pTwoTensorPos);
	for (int i = 0; i < count; ++i) {
		SafeDeleteArray(pTwoTensorBuffer[i]);
	}
	SafeDeleteArray(pTwoTensorBuffer);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDWI = new CDWI();
	pDWI->OpenDWIFile(pDWIPathname);
	pDTI = new CDTI();
	pDTI->OpenDTIFile(pDTIPathname);
	pModelSelection = new CScalar();
	pModelSelection->OpenScalarFile(pModelSelectionPathname);

	CreateMLSSamplePoints();
	RegularizedTwoTensorFitting();

	SafeDelete(pDWI);
	SafeDelete(pDTI);
	SafeDelete(pModelSelection);

	return 0;
}
